# 采样用工具，可进行标定/训练素材的采样，使用前(主要是标定时)需先确认左右相机是否对应程序中的左和右
# 运行该程序时请注意是否已经创建了SAVE_PATH对应文件夹
import cv2
from stereo_config_200 import stereoCamera

SAVE_PATH = "C:\\Users\\fsf\\Desktop\\imgs\\"  # 自行修改并创建imgs文件夹,若运行标定采样程序则在其下再创建l和r两个文件夹

# 标定用拍照，按s拍，存成0_l/0_r, 1_l/1_r, 2_l/2_r, ......
def for_calibration(start_index = 0):
    print("Image Capture For Calibration, 's' to shot, 'q' to quit")
    cap = cv2.VideoCapture(1)
    cap2 = cv2.VideoCapture(2)
    i = start_index
    while(True):
        ret, frame = cap.read()
        ret2, frame2 = cap2.read()
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)  # matlab2020双目标定只支持输入灰度图
        frame2 = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)
        k = cv2.waitKey(1)
        if k == ord('q'):
            break
        elif k == ord('s'):
            cv2.imwrite(SAVE_PATH + 'l\\' + str(i) + '_l.jpg', frame)
            cv2.imwrite(SAVE_PATH + 'r\\' + str(i) + '_r.jpg', frame2)
            i += 1
        cv2.imshow("left", frame)
        cv2.imshow("right", frame2)
    cap.release()
    cap2.release()
    cv2.destroyAllWindows()

# 训练用拍照，按s拍，从start_index开始存成0, 1, 2, 3, ......
def for_training(start_index = 0):
    print("Image Capture For Training, 's' to shot, 'q' to quit")
    config = stereoCamera()
    cap = cv2.VideoCapture(1)
    cap2 = cv2.VideoCapture(2)
    i = start_index
    while(True):
        ret, frame = cap.read()
        ret2, frame2 = cap2.read()
        # 训练用照片就做畸变校正
        frame, frame2 = config.Rectify(frame, frame2)
        k = cv2.waitKey(1)
        if k == ord('q'):
            break
        elif k == ord('s'):
            # 因定位部分对左右两张图片的识别结果都有很大依赖，而左右相机存在一定色差，故左右两张图都存
            cv2.imwrite(SAVE_PATH + str(i) + '.jpg', frame)
            i += 1
            cv2.imwrite(SAVE_PATH + str(i) + '.jpg', frame2)
            i += 1
        cv2.imshow("left", frame)
        cv2.imshow("right", frame2)
    cap.release()
    cap2.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    #for_calibration(start_index = 0)
    for_training(start_index = 0)